#! /usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author :   Liang Kang
@Contact:   gangkanli1219@163.com
@Time   :   2018/7/31 13:42
@Desc   :
"""
import hashlib
import os

import cv2
import numpy as np
import tensorflow as tf
from dltools.data import BaseRecordGenerator
from dltools.utils.io import read_voc_xml
from object_detection.utils import dataset_util


def _create_example(objects, area,
                    image_path, xml_path, shape,
                    img_format='JPEG'):
    """
    创建一个 tensorflow example
    Parameters
    ----------
    objects
    area
    image_path
    xml_path
    shape
    img_format

    Returns
    -------

    """
    boxes, labels, cls_name = [], [], []
    for ob in objects:
        box = ob['box']
        if (box[3] - box[1]) * (box[2] - box[0]) < area:
            continue
        labels.append(ob['label'])
        boxes.append(ob['box'])
        cls_name.append('shelf'.encode('utf8'))
    if len(labels) > 0:
        labels = np.ones((len(labels),), dtype=np.int32)
        boxes = np.asarray(boxes) / np.array([shape[:2] * 2])
        boxes = np.clip(boxes, 0.0, 1.0)
        axises = np.split(boxes, 4, 1)
        axises = list(map(lambda x: np.squeeze(x).tolist(), axises))
        ymin, xmin, ymax, xmax = axises
        if labels.size == 1:
            ymin, xmin, ymax, xmax = [ymin], [xmin], [ymax], [xmax]
        labels = labels.tolist()
    else:
        return None
    with tf.gfile.GFile(image_path, 'rb') as fid:
        encoded_jpg = fid.read()
    key = hashlib.sha256(encoded_jpg).hexdigest()
    truncated, poses, difficult_obj = [], [], []
    example = tf.train.Example(
        features=tf.train.Features(feature={
            'image/height':
                dataset_util.int64_feature(shape[0]),
            'image/width':
                dataset_util.int64_feature(shape[1]),
            'image/filename':
                dataset_util.bytes_feature(
                    os.path.basename(image_path).encode('utf8')),
            'image/source_id':
                dataset_util.bytes_feature(
                    os.path.basename(xml_path).encode('utf8')),
            'image/key/sha256':
                dataset_util.bytes_feature(key.encode('utf8')),
            'image/encoded':
                dataset_util.bytes_feature(encoded_jpg),
            'image/format':
                dataset_util.bytes_feature(img_format.encode('utf8')),
            'image/object/bbox/xmin':
                dataset_util.float_list_feature(xmin),
            'image/object/bbox/xmax':
                dataset_util.float_list_feature(xmax),
            'image/object/bbox/ymin':
                dataset_util.float_list_feature(ymin),
            'image/object/bbox/ymax':
                dataset_util.float_list_feature(ymax),
            'image/object/class/text':
                dataset_util.bytes_list_feature(cls_name),
            'image/object/class/label':
                dataset_util.int64_list_feature(labels),
            'image/object/difficult':
                dataset_util.int64_list_feature(difficult_obj),
            'image/object/truncated':
                dataset_util.int64_list_feature(truncated),
            'image/object/view':
                dataset_util.bytes_list_feature(poses),
        }))
    return example.SerializeToString()


class ShelfRecordGenerator(BaseRecordGenerator):

    def __init__(self, data, output, area=1000, *args, **kwargs):
        super(ShelfRecordGenerator, self).__init__(data, output, *args,
                                                   **kwargs)
        self.area = area

    def _encode_data(self):
        image_path = self._buf_data['raw']['image']
        xml_path = self._buf_data['raw']['xml']
        image = cv2.imread(image_path)
        shape = image.shape
        image_info, flag = read_voc_xml(xml_path, shape)
        if flag:
            self._buf_data['data'] = None
            return
        self._buf_data['data'] = _create_example(
            image_info['objects'], self.area, image_path, xml_path, shape)

    def _write_data(self, writer):
        example = self._buf_data['data']
        if example is None:
            return
        writer.write(example)
